
package cartagows.wsframework.wssecurity.core;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Vector;

import org.apache.axiom.om.OMElement;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ws.security.WSConstants;
import org.apache.ws.security.WSEncryptionPart;
import org.apache.ws.security.WSSecurityException;
import org.apache.ws.security.components.crypto.Crypto;
import org.apache.ws.security.components.crypto.CryptoFactory;
import org.apache.ws.security.conversation.ConversationException;
import org.apache.ws.security.message.WSSecDKEncrypt;
import org.apache.ws.security.message.WSSecDKSign;
import org.apache.ws.security.message.WSSecEncrypt;
import org.apache.ws.security.message.WSSecEncryptedKey;
import org.apache.ws.security.message.WSSecSignature;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import cartagows.wsframework.wssecurity.policy.PolicyData;
import cartagows.wsframework.wssecurity.policy.SPConstants;
import cartagows.wsframework.wssecurity.policy.model.AlgorithmSuite;
import cartagows.wsframework.wssecurity.policy.model.SupportingToken;
import cartagows.wsframework.wssecurity.policy.model.Token;
import cartagows.wsframework.wssecurity.utils.SecurityUtil;

public class AsymmetricBindingBuilder extends BindingBuilder {

	private static Log log = LogFactory.getLog(AsymmetricBindingBuilder.class);

	private Token sigToken;
	private WSSecSignature sig;
	private WSSecEncryptedKey encrKey;
	private String encryptedKeyId;
	private byte[] encryptedKeyValue;
	private Vector signatureValues = new Vector();
	private Element encrTokenElement;
	private Element sigDKTElement;
	private Element encrDKTElement;
	private Vector sigParts = new Vector();
	private Element signatureElement; 

	public AsymmetricBindingBuilder(){
	}

	public void build(MessageData md) throws SecurityException {
		log.debug("AsymmetricBindingBuilder build invoked");

		PolicyData pd = md.getPolicyData();
		if (pd.isIncludeTimestamp()) {
			this.addTimestamp(md);
		}

		if (SPConstants.ENCRYPT_BEFORE_SIGNING.equals(pd.getProtectionOrder())) {
			this.doEncryptBeforeSig(md);
		} else {
			this.doSignBeforeEncrypt(md);
		} 

		log.debug("AsymmetricBindingBuilder build invoked : DONE");
	}

	private void doEncryptBeforeSig(MessageData md) throws SecurityException {

		PolicyData pd = md.getPolicyData();
		Document doc = md.getDocument();

		/*
		 * We need to hold on to these two element to use them as refence in the
		 * case of encypting the signature
		 */
		Element encrDKTokenElem = null;
		WSSecEncrypt encr = null;
		Element refList = null;
		WSSecDKEncrypt dkEncr = null;

		/*
		 * We MUST use keys derived from the same token
		 */
		Token encryptionToken = null;
		if(md.isInitiator()) {
			encryptionToken = pd.getRecipientToken();
		} else {
			encryptionToken = pd.getInitiatorToken();
		}
		Vector encrParts = SecurityUtil.getEncryptedParts(md);

		//Signed parts are determined before encryption because encrypted signed  headers
		//will not be included otherwise
		this.sigParts = SecurityUtil.getSignedParts(md);

		if(encryptionToken == null && encrParts.size() > 0) {
			throw new SecurityException("encryptionTokenMissing");
		}

		if (encryptionToken != null && encrParts.size() > 0) {

			if (encryptionToken.isDerivedKeys()) {
				try {
					this.setupEncryptedKey(md, encryptionToken);
					// Create the DK encryption builder
					dkEncr = new WSSecDKEncrypt();
					dkEncr.setParts(encrParts);
					dkEncr.setExternalKey(this.encryptedKeyValue, this.encryptedKeyId);
					dkEncr.setDerivedKeyLength(pd.getAlgorithmSuite().getEncryptionDerivedKeyLength()/8);
					dkEncr.prepare(doc);

					// Get and add the DKT element
					this.encrDKTElement = dkEncr.getdktElement();
					encrDKTokenElem = SecurityUtil.appendChildToSecHeader(md, this.encrDKTElement);

					refList = dkEncr.encryptForExternalRef(null, encrParts);

				} catch (WSSecurityException e) {
					throw new SecurityException("errorCreatingEncryptedKey");
				} catch (ConversationException e) {
					throw new SecurityException("errorInDKEncr");
				}
			} else {
				try {
					encr = new WSSecEncrypt();
					encr.setParts(encrParts);
					encr.setDocument(doc);
					SecurityUtil.setEncryptionUser(md, encr);
					encr.setSymmetricEncAlgorithm(pd.getAlgorithmSuite().getEncryption());
					SecurityUtil.setKeyIdentifierType(pd,encr, encryptionToken);
					encr.setKeyEncAlgo(pd.getAlgorithmSuite().getAsymmetricKeyWrap());
					Crypto crypto = CryptoFactory.getInstance(md.getCryptoEncProperties());
					encr.prepare(doc, crypto);

					Element bstElem = encr.getBinarySecurityTokenElement();
					if (bstElem != null) {
						SecurityUtil.appendChildToSecHeader(md, bstElem);
					}

					this.encrTokenElement = encr.getEncryptedKeyElement();
					this.encrTokenElement = SecurityUtil.appendChildToSecHeader(md,
							encrTokenElement);

					refList = encr.encryptForExternalRef(null, encrParts);

				} catch (WSSecurityException e) {
					throw new SecurityException("errorInEncryption");
				}
			}

			SecurityUtil.appendChildToSecHeader(md, refList);

			this.setInsertionLocation(encrTokenElement);

			SecurityUtil.handleEncryptedSignedHeaders(encrParts, this.sigParts, doc);

			HashMap sigSuppTokMap = null;
			HashMap endSuppTokMap = null;
			HashMap sgndEndSuppTokMap = null;
			HashMap sgndEncSuppTokMap = null;
			HashMap endEncSuppTokMap = null;
			HashMap sgndEndEncSuppTokMap = null;

			if(this.timestampElement != null){
				sigParts.add(new WSEncryptionPart(SecurityUtil
						.addWsuIdToElement((OMElement) this.timestampElement)));
			}

			if (md.isInitiator()) {

				// Now add the supporting tokens
				SupportingToken sgndSuppTokens = pd.getSignedSupportingTokens();
				sigSuppTokMap = this.handleSupportingTokens(md, sgndSuppTokens);           

				SupportingToken endSuppTokens = pd.getEndorsingSupportingTokens();
				endSuppTokMap = this.handleSupportingTokens(md, endSuppTokens);

				SupportingToken sgndEndSuppTokens = pd.getSignedEndorsingSupportingTokens();           
				sgndEndSuppTokMap = this.handleSupportingTokens(md, sgndEndSuppTokens);

				SupportingToken sgndEncryptedSuppTokens = pd.getSignedEncryptedSupportingTokens();
				sgndEncSuppTokMap = this.handleSupportingTokens(md, sgndEncryptedSuppTokens);

				SupportingToken endorsingEncryptedSuppTokens = pd.getEndorsingEncryptedSupportingTokens();
				endEncSuppTokMap = this.handleSupportingTokens(md, endorsingEncryptedSuppTokens);

				SupportingToken sgndEndEncSuppTokens = pd.getSignedEndorsingEncryptedSupportingTokens();           
				sgndEndEncSuppTokMap = this.handleSupportingTokens(md, sgndEndEncSuppTokens);

				SupportingToken supportingToks = pd.getSupportingTokens();
				this.handleSupportingTokens(md, supportingToks);

				SupportingToken encryptedSupportingToks = pd.getEncryptedSupportingTokens();
				this.handleSupportingTokens(md, encryptedSupportingToks);

				//Setup signature parts
				sigParts = addSignatureParts(sigSuppTokMap, sigParts);
				sigParts = addSignatureParts(sgndEncSuppTokMap, sigParts);
				sigParts = addSignatureParts(sgndEndSuppTokMap, sigParts);
				sigParts = addSignatureParts(sgndEndEncSuppTokMap, sigParts);

			} else {
                addSignatureConfirmation(md, sigParts);
            }

			if(( sigParts.size() > 0 &&
					md.isInitiator() && pd.getInitiatorToken() != null) || 
					(!md.isInitiator() && pd.getRecipientToken() != null)) {
				this.doSignature(md);
			}

			if (md.isInitiator()) {
                
                endSuppTokMap.putAll(endEncSuppTokMap);
                // Do endorsed signatures
                Vector endSigVals = this.doEndorsedSignatures(md,
                        endSuppTokMap);
                for (Iterator iter = endSigVals.iterator(); iter.hasNext();) {
                    signatureValues.add(iter.next());
                }

                sgndEndSuppTokMap.putAll(sgndEndEncSuppTokMap);
                // Do signed endorsing signatures
                Vector sigEndSigVals = this.doEndorsedSignatures(md,
                        sgndEndSuppTokMap);
                for (Iterator iter = sigEndSigVals.iterator(); iter.hasNext();) {
                    signatureValues.add(iter.next());
                }
            }
			
			// Check for signature protection
			if (pd.isSignatureProtection() && this.mainSigId != null) {

				Vector secondEncrParts = new Vector();

				// Now encrypt the signature using the above token
				secondEncrParts.add(new WSEncryptionPart(this.mainSigId,
				"Element"));

				if(md.isInitiator()) {
					for (int i = 0 ; i < encryptedTokensIdList.size(); i++) {
						secondEncrParts.add(new WSEncryptionPart((String)encryptedTokensIdList.get(i),"Element"));
					}
				}

				Element secondRefList = null;

				if (encryptionToken.isDerivedKeys()) {
					try {

						secondRefList = dkEncr.encryptForExternalRef(null,
								secondEncrParts);
						SecurityUtil.insertSiblingAfter(md, encrDKTokenElem,
								secondRefList);

					} catch (WSSecurityException e) {
						throw new SecurityException("errorCreatingEncryptedKey");
					}
				} else {
					try {
						// Encrypt, get hold of the ref list and add it
						secondRefList = encr.encryptForExternalRef(null,
								secondEncrParts);

						// Insert the ref list after the encrypted key elem
						this.setInsertionLocation(SecurityUtil
								.insertSiblingAfter(md, encrTokenElement,
										secondRefList));
					} catch (WSSecurityException e) {
						throw new SecurityException("errorInEncryption");
					}
				}

			}
		}

	}

	private void doSignBeforeEncrypt(MessageData md) throws SecurityException {

		PolicyData pd = md.getPolicyData();
		Document doc = md.getDocument();

		HashMap sigSuppTokMap = null;
		HashMap endSuppTokMap = null;
		HashMap sgndEndSuppTokMap = null;
		HashMap sgndEncSuppTokMap = null;
		HashMap endEncSuppTokMap = null;
		HashMap sgndEndEncSuppTokMap = null;

		sigParts = SecurityUtil.getSignedParts(md);
		
		//Add timestamp
		if(this.timestampElement != null){
			sigParts.add(new WSEncryptionPart(SecurityUtil
					.addWsuIdToElement((OMElement) this.timestampElement)));
		}else{
			this.setInsertionLocation(null);
		}

		if (md.isInitiator()) {

			//      Now add the supporting tokens
			SupportingToken sgndSuppTokens = pd.getSignedSupportingTokens();
			sigSuppTokMap = this.handleSupportingTokens(md, sgndSuppTokens);           

			SupportingToken endSuppTokens = pd.getEndorsingSupportingTokens();
			endSuppTokMap = this.handleSupportingTokens(md, endSuppTokens);

			SupportingToken sgndEndSuppTokens = pd.getSignedEndorsingSupportingTokens();           
			sgndEndSuppTokMap = this.handleSupportingTokens(md, sgndEndSuppTokens);

			SupportingToken sgndEncryptedSuppTokens = pd.getSignedEncryptedSupportingTokens();
			sgndEncSuppTokMap = this.handleSupportingTokens(md, sgndEncryptedSuppTokens);

			SupportingToken endorsingEncryptedSuppTokens = pd.getEndorsingEncryptedSupportingTokens();
			endEncSuppTokMap = this.handleSupportingTokens(md, endorsingEncryptedSuppTokens);

			SupportingToken sgndEndEncSuppTokens = pd.getSignedEndorsingEncryptedSupportingTokens();           
			sgndEndEncSuppTokMap = this.handleSupportingTokens(md, sgndEndEncSuppTokens);

			SupportingToken supportingToks = pd.getSupportingTokens();
			this.handleSupportingTokens(md, supportingToks);

			SupportingToken encryptedSupportingToks = pd.getEncryptedSupportingTokens();
			this.handleSupportingTokens(md, encryptedSupportingToks);

			//Setup signature parts
			sigParts = addSignatureParts(sigSuppTokMap, sigParts);
			sigParts = addSignatureParts(sgndEncSuppTokMap, sigParts);
			sigParts = addSignatureParts(sgndEndSuppTokMap, sigParts);
			sigParts = addSignatureParts(sgndEndEncSuppTokMap, sigParts);

		} else {
            addSignatureConfirmation(md, sigParts);
        }

		if( sigParts.size() > 0 && 
				((md.isInitiator() && pd.getInitiatorToken() != null) || 
						(!md.isInitiator() && pd.getRecipientToken() != null))) {
			// Do signature
			this.doSignature(md);
		}

		//Do endorsed signature

        if (md.isInitiator()) {
            
            // Adding the endorsing encrypted supporting tokens to endorsing supporting tokens
            endSuppTokMap.putAll(endEncSuppTokMap);
            // Do endorsed signatures
            Vector endSigVals = this.doEndorsedSignatures(md,
                    endSuppTokMap);
            for (Iterator iter = endSigVals.iterator(); iter.hasNext();) {
                signatureValues.add(iter.next());
            }

            //Adding the signed endorsed encrypted tokens to signed endorsed supporting tokens
            sgndEndSuppTokMap.putAll(sgndEndEncSuppTokMap);
            // Do signed endorsing signatures
            Vector sigEndSigVals = this.doEndorsedSignatures(md,
                    sgndEndSuppTokMap);
            for (Iterator iter = sigEndSigVals.iterator(); iter.hasNext();) {
                signatureValues.add(iter.next());
            }
        }
		
		Vector encrParts = SecurityUtil.getEncryptedParts(md);

		//Check for signature protection
		if(pd.isSignatureProtection() && this.mainSigId != null) {
			encrParts.add(new WSEncryptionPart(SecurityUtil.addWsuIdToElement((OMElement)this.signatureElement), "Element"));
		}

		if(md.isInitiator()) {
			for (int i = 0 ; i < encryptedTokensIdList.size(); i++) {
				encrParts.add(new WSEncryptionPart((String)encryptedTokensIdList.get(i),"Element"));
			}
		}

		//Do encryption
		Token encrToken = pd.getRecipientToken();
		if(encrToken != null && encrParts.size() > 0) {
			Element refList = null;
			AlgorithmSuite algorithmSuite = pd.getAlgorithmSuite();
			if(encrToken.isDerivedKeys()) {

				try {
					WSSecDKEncrypt dkEncr = new WSSecDKEncrypt();

					if(this.encrKey == null) {
						this.setupEncryptedKey(md, encrToken);
					}

					dkEncr.setExternalKey(this.encryptedKeyValue, this.encryptedKeyId);
					dkEncr.setCustomValueType(WSConstants.SOAPMESSAGE_NS11 + "#"
							+ WSConstants.ENC_KEY_VALUE_TYPE);
					dkEncr.setSymmetricEncAlgorithm(algorithmSuite.getEncryption());
					dkEncr.setDerivedKeyLength(algorithmSuite.getEncryptionDerivedKeyLength()/8);
					dkEncr.prepare(doc);


					if(this.encrTokenElement != null) {
						this.encrDKTElement = SecurityUtil.insertSiblingAfter(
								md, this.encrTokenElement, dkEncr.getdktElement());
					} else {
						this.encrDKTElement = SecurityUtil.insertSiblingBefore(
								md, this.sigDKTElement, dkEncr.getdktElement());
					}

					refList = dkEncr.encryptForExternalRef(null, encrParts);

					SecurityUtil.insertSiblingAfter(md, 
							this.encrDKTElement, 
							refList);

				} catch (WSSecurityException e) {
					throw new SecurityException("errorInDKEncr");
				} catch (ConversationException e) {
					throw new SecurityException("errorInDKEncr");
				}
			} else {
				try {

					WSSecEncrypt encr = new WSSecEncrypt();

					SecurityUtil.setKeyIdentifierType(pd, encr, encrToken);

					encr.setDocument(doc);
					SecurityUtil.setEncryptionUser(md, encr);
					encr.setSymmetricEncAlgorithm(algorithmSuite.getEncryption());
					encr.setKeyEncAlgo(algorithmSuite.getAsymmetricKeyWrap());
					Crypto crypto = CryptoFactory.getInstance(md.getCryptoEncProperties());
					encr.prepare(doc, crypto);

					if(this.timestampElement != null){
						this.setInsertionLocation(this.timestampElement);
					}else{
						this.setInsertionLocation(null);
					}

					if(encr.getBSTTokenId() != null) {
						this.setInsertionLocation(SecurityUtil
								.insertSiblingAfterOrPrepend(md,
										this.getInsertionLocation(),
										encr.getBinarySecurityTokenElement()));
					}


					Element encryptedKeyElement = encr.getEncryptedKeyElement();

					//Encrypt, get hold of the ref list and add it
					refList = encr.encryptForInternalRef(null, encrParts);

					//Add internal refs
					encryptedKeyElement.appendChild(refList);

					this.setInsertionLocation(SecurityUtil
							.insertSiblingAfterOrPrepend(md,
									this.getInsertionLocation(),
									encryptedKeyElement)); 
				} catch (WSSecurityException e) {
					throw new SecurityException("errorInEncryption");
				}    
			}
		}

	}

	private void doSignature(MessageData md) throws SecurityException {

		PolicyData pd = md.getPolicyData();
		Document doc = md.getDocument();

		if(md.isInitiator()) {
			sigToken = pd.getInitiatorToken();
		} else {
			sigToken = pd.getRecipientToken();
		}

		if (sigToken.isDerivedKeys()) {
			// Set up the encrypted key to use
			if(this.encrKey == null) {
				setupEncryptedKey(md, sigToken);
			}

			WSSecDKSign dkSign = new WSSecDKSign();
			dkSign.setExternalKey(this.encryptedKeyValue, this.encryptedKeyId);

			// Set the algo info
			dkSign.setSignatureAlgorithm(pd.getAlgorithmSuite()
					.getSymmetricSignature());
			dkSign.setDerivedKeyLength(pd.getAlgorithmSuite()
					.getSignatureDerivedKeyLength() / 8);
			dkSign.setCustomValueType(WSConstants.SOAPMESSAGE_NS11 + "#"
					+ WSConstants.ENC_KEY_VALUE_TYPE);
			try {
				dkSign.prepare(doc, md.getSecHeader());

				if (pd.isTokenProtection()) {
					sigParts.add(new WSEncryptionPart(encrKey.getId()));
				}

				dkSign.setParts(sigParts);

				dkSign.addReferencesToSign(sigParts, md.getSecHeader());

				// Do signature
				dkSign.computeSignature();

				;
				// Add elements to header
				this.sigDKTElement = SecurityUtil.insertSiblingAfter(md,
						this.getInsertionLocation(), dkSign.getdktElement());
				this.setInsertionLocation(this.sigDKTElement);

				this.setInsertionLocation(SecurityUtil.insertSiblingAfter(md,
						this.getInsertionLocation(), dkSign
						.getSignatureElement()));

				this.mainSigId = SecurityUtil
				.addWsuIdToElement((OMElement) dkSign
						.getSignatureElement());

				signatureValues.add(dkSign.getSignatureValue());

				signatureElement = dkSign.getSignatureElement();
			} catch (WSSecurityException e) {
				throw new SecurityException("errorInDerivedKeyTokenSignature");
			} catch (ConversationException e) {
				throw new SecurityException("errorInDerivedKeyTokenSignature");
			}

		} else {
			sig = this.getSignatureBuilder(md, sigToken);
			Element bstElem = sig.getBinarySecurityTokenElement();
			if(bstElem != null) {
				bstElem = SecurityUtil.insertSiblingAfter(md, this
						.getInsertionLocation(), bstElem);
				this.setInsertionLocation(bstElem);
			}

			if (md.getPolicyData().isTokenProtection()
					&& sig.getBSTTokenId() != null) {
				sigParts.add(new WSEncryptionPart(sig.getBSTTokenId()));
			}

			try {
				sig.addReferencesToSign(sigParts, md.getSecHeader());
				sig.computeSignature();

				signatureElement = sig.getSignatureElement();

				this.setInsertionLocation(SecurityUtil.insertSiblingAfter(
						md, this.getInsertionLocation(), signatureElement));

				this.mainSigId = SecurityUtil.addWsuIdToElement((OMElement) signatureElement);
			} catch (WSSecurityException e) {
				throw new SecurityException("errorInSignatureWithX509Token");
			}
			signatureValues.add(sig.getSignatureValue());
		}

	}

	private void setupEncryptedKey(MessageData md, Token token) throws SecurityException {
		if(!md.isInitiator() && token.isDerivedKeys()) {

			//If we already have them, simply return
            if(this.encryptedKeyId != null && this.encryptedKeyValue != null) {
                return;
            }
            
            //Use the secret from the incoming EncryptedKey element
            Object resultsObj = md.getResults();
            if(resultsObj != null) {
                encryptedKeyId = SecurityUtil.getRequestEncryptedKeyId((Vector)resultsObj);
                encryptedKeyValue = SecurityUtil.getRequestEncryptedKeyValue((Vector)resultsObj);
                
                //In the case where we don't have the EncryptedKey in the 
                //request, for the control to have reached this state,
                //the scenario MUST be a case where this is the response
                //message by a listener created for an async client
                //Therefor we will create a new EncryptedKey
                if(encryptedKeyId == null && encryptedKeyValue == null) {
                    createEncryptedKey(md, token);
                }
            } else {
                throw new SecurityException("noSecurityResults");
            }
    } else {
        createEncryptedKey(md, token);
    }

	}

	private void createEncryptedKey(MessageData md, Token token) throws SecurityException {
		//Set up the encrypted key to use
		encrKey = this.getEncryptedKeyBuilder(md, token);

		Element bstElem = encrKey.getBinarySecurityTokenElement();
		if (bstElem != null) {
			// If a BST is available then use it
			SecurityUtil.appendChildToSecHeader(md, bstElem);
		}

		// Add the EncryptedKey
		encrTokenElement = encrKey.getEncryptedKeyElement();
		this.encrTokenElement = SecurityUtil.appendChildToSecHeader(md,
				encrTokenElement);
		encryptedKeyValue = encrKey.getEphemeralKey();
		encryptedKeyId = encrKey.getId();
	}
}
